from milvus import Milvus, IndexType, MetricType, Status
import milvus
import random

class Milvus_op():
    def __init__(self, uri, conn_server=True,print_enable=False):

        if conn_server:
            self.connect_server(uri)

        self.print_enable = print_enable

    def connect_server(self, uri):
        self.milvus = Milvus(uri=uri)

    def set_op(self, op='c', param=None):

        assert (op != None)
        assert (param != None)
        assert (param['collection_name'])  # 集合名

        if op == 'c':
            # 创建集合
            assert (param['dimension'])  # 向量维度
            assert (param['index_file_size'])  # 索引文件大小
            assert (param['metric_type'])  # 比较向量方式
            status = self.milvus.create_collection(param)

            # 默认自动创建索引
            self.index_op(op='c',collection_name=param['collection_name'],param={'nlist':16384})


        elif op == 'd':
            assert (param['collection_name'])  # 集合名
            status = self.milvus.drop_collection(collection_name=param['collection_name'])
            # 默认自动删除
            self.index_op(op='d', collection_name=param['collection_name'])

        if self.print_enable:
            print("Milvus:collection operation,", param['collection_name'], status)

    # 分区操作
    def partition_op(self, op='c', set_name=None, partition_name=None):
        assert (op != None)
        assert (set_name != None)
        assert (partition_name != None)
        if op == 'c':
            # 创建分区
            status = self.milvus.create_partition(set_name, partition_name)

        elif op == 'd':
            # 删除分区
            status = self.milvus.drop_partition(collection_name=set_name, partition_tag=partition_name)
        if self.print_enable:
            print(f"Milvus:partition operation,set: {set_name} ,partition: {partition_name}", status)

    # 向量操作
    def vectors_op(self, op, collection_name, vectors, partitions=None, ids=None):
        assert (op != None)
        assert (collection_name != None)
        assert (vectors != None)

        if op == 'i':
            if ids == None and partitions == None:

                status, ret_id = self.milvus.insert(collection_name=collection_name, records=vectors)
            elif ids != None and partitions == None:

                status, ret_id = self.milvus.insert(collection_name=collection_name, records=vectors, ids=ids)
            elif ids != None and partitions != None:

                status, ret_id = self.milvus.insert(collection_name=collection_name, partition_tag=partitions,
                                                    records=vectors, ids=ids)
            elif ids == None and partitions != None:

                status, ret_id = self.milvus.insert(collection_name=collection_name, partition_tag=partitions,
                                                    records=vectors)
            if self.print_enable:
                print(f"Milvus:Insert vectors",status)
            return ret_id
        elif op == 'd':
            assert (ids != None)
            status = self.milvus.delete_entity_by_id(collection_name=collection_name, id_array=ids)
            self.milvus.flush(collection_name)
            if self.print_enable:
                print(f"Milvus:Delete vectors {ids}", status)

    def index_op(self, op, collection_name, param, indexType=IndexType.IVF_FLAT):

        assert (op != None)
        assert (collection_name != None)
        assert (param != None)
        if op == 'c':
            assert (param['nlist'])
            self.milvus.create_index(collection_name, indexType, param)
        elif op == 'd':
            self.milvus.drop_index(collection_name)
            self.milvus.flush(collection_name)

    def search_vectors(self, search_param, collection_name, q_records, top_k=2, partition_tags=None):
        assert (search_param != None)
        assert (collection_name != None)
        assert (q_records != None)

        if partition_tags != None:
            status, results = self.milvus.search(collection_name=collection_name, query_records=q_records, top_k=top_k,
                                                 partition_tags=partition_tags,
                                                 params=search_param)
        else:
            status, results = self.milvus.search(collection_name=collection_name, query_records=q_records, top_k=top_k,
                                                 params=search_param)
        if self.print_enable:
            print("search vectors ok", status)
        return results

    def flush(self,collection_name):
        status = self.milvus.flush(collection_name_array=collection_name)
        if self.print_enable:
            print(f"flush {collection_name} ok,{status}")
    def close_client(self):
        self.milvus.close()


# 连接服务端
# def connect_server(uri):
#     milvus = Milvus(uri='tcp://localhost:19530')
#     return milvus
# milvus的使用


if __name__ == '__main__':
    mymilvus = Milvus_op('tcp://192.168.1.56:19530',print_enable=True)
    # 欧式距离
    param = {'collection_name': 'test02', 'dimension': 256, 'index_file_size': 1024, 'metric_type': MetricType.L2}
    mymilvus.set_op('c',param)
    mymilvus.partition_op('c','test02','partion')

    # 查询向量
    vectors = [[random.random() for _ in range(256)] for _ in range(20)]
    rets = mymilvus.search_vectors(search_param = {'nprobe':16},collection_name='test02',q_records=vectors)
    print(rets)

    # ids = [1]*20
    # ids = mymilvus.vectors_op(op='i',collection_name='test02',vectors=vectors)
    # print(ids)

    mymilvus.close_client()
    # print(mymilvus.milvus)
